#pragma once

#include "EK.h"

using namespace std;

struct Center
{
    pair<double,double> loc;
    bool workerFlag = false;
    bool taskFlag = false;
};
struct dm_node
{
    double dmScore;
    int index1;
    dm_node(double maScore,int index1):
        dmScore(maScore),index1(index1){}
    //定义从小到大排序    
    bool operator<(const dm_node& o) const {
		return o.dmScore < dmScore;
	}
};

double calEuclidean(pair<double,double>& worker1,pair<double,double>& worker2,pair<double,double>& task1,pair<double,double>& task2){
    double workerDistance=0,taskDistance=0;
    
    // double mdist = MahalanobisDistance2(worker1,worker2,task1,task2);
    double mdist = -1;
    if(mdist !=-1) {
        return mdist;
    }
    else{
        // cout<<"gg"<<endl;
        workerDistance = sqrt((worker1.first-worker2.first)*(worker1.first-worker2.first)+(worker1.second-worker2.second)*(worker1.second-worker2.second));
        taskDistance = sqrt((task1.first-task2.first)*(task1.first-task2.first)+(task1.second-task2.second)*(task1.second-task2.second));
        // double wtDistance1 = sqrt((worker1.first-task1.first)*(worker1.first-task1.first)+(worker1.second-task1.second)*(worker1.second-task1.second));
        // double wtDistance2 = sqrt((worker2.first-task2.first)*(worker2.first-task2.first)+(worker2.second-task2.second)*(worker2.second-task2.second));
        // double wtDistance1 = sqrt((worker1.first-task2.first)*(worker1.first-task2.first)+(worker1.second-task2.second)*(worker1.second-task2.second));
        // double wtDistance2 = sqrt((worker2.first-task1.first)*(worker2.first-task1.first)+(worker2.second-task1.second)*(worker2.second-task1.second));
    }
    return workerDistance+taskDistance;
}


void initCenter(int&MSize,int&d,vector<Center>&muWorkers,vector<Center>&muTasks){
    //先随机选择一个中心点
    Center cTemp;
    int initCenter = rand()%(MSize);
    cTemp.loc = perturbWorkers[M0[initCenter].first].loc;
    muWorkers.push_back(cTemp);
    cTemp.loc = perturbTasks[M0[initCenter].second].loc;
    muTasks.push_back(cTemp);
    int muSize = muTasks.size();
    vector<vector<double>> Madistance(MSize,vector<double>(MSize,-1));//马氏距离保存矩阵，防止重复计算带来时间开销
    vector<int> centerM0;//记录匹配对中M0中哪一个为中心
    centerM0.push_back(initCenter);
    for(int i=0;i<MSize;i++){
        Madistance[i][i] = 0;
        for(int j=i;j<MSize;j++){
            double madis = calEuclidean(perturbWorkers[M0[i].first].loc,perturbWorkers[M0[j].first].loc,perturbTasks[M0[i].second].loc,perturbTasks[M0[j].second].loc); 
            // cout<<madis<<endl;
            Madistance[i][j] = madis;
            Madistance[j][i] = madis;//矩阵是对称的
        }
    }
    //获取初始的中心点
    while (muSize < d)
    {
        double maxDistance=-1,minDistance,minDTemp;
        int centerTemp = -1;
        pair<int,int> sample;
        for(int i=0;i<MSize;i++){
            sample = M0[i];
            minDistance=10e10;
            for(int j=0;j<muSize;j++){
                // minDTemp = calEuclidean(perturbWorkers[sample.first].loc,muWorkers[j].loc,perturbTasks[sample.second].loc,muTasks[j].loc);
                minDTemp = Madistance[i][centerM0[j]];
                if(minDTemp < minDistance){
                    minDistance = minDTemp;
                }
            }
            if(minDistance > maxDistance){
                maxDistance = minDistance;
                centerTemp = i;
            }
        }
        cTemp.loc = perturbWorkers[M0[centerTemp].first].loc;
        centerM0.push_back(centerTemp);//记录中心
        muWorkers.push_back(cTemp);
        cTemp.loc = perturbTasks[M0[centerTemp].second].loc;
        muTasks.push_back(cTemp);
        muSize = muTasks.size();
    }
}

int calMeans(int&MSize,int&d,vector<Center>&muWorkers,vector<Center>&muTasks,vector<vector<pair<int,int>>>& cluster){
    for(int i=0;i<d;i++){
        cluster[i].clear();
    }
    // cluster.clear();
    // cluster.resize(d,vector<pair<int,int>>());
    //分类为哪一个簇
    pair<int,int> xi;
    // double min,minTemp;
    int clusterIndex=-1;
    vector<vector<double>> Centerdist(MSize,vector<double>(d,-1));//每一个点与簇中心的距离
    // vector<vector<dm_node>>Centerdist2;
    // vector<dm_node> disHeap;
    for(int i=0;i<MSize;i++){
        xi = M0[i];
        // disHeap.clear();
        for(int j=0;j<d;j++){
            Centerdist[i][j] = calEuclidean(perturbWorkers[xi.first].loc,muWorkers[j].loc,perturbTasks[xi.second].loc,muTasks[j].loc);
            // double dm = calEuclidean(perturbWorkers[xi.first].loc,muWorkers[j].loc,perturbTasks[xi.second].loc,muTasks[j].loc);
            // disHeap.push_back(dm_node(dm,j));
            // push_heap(disHeap.begin(),disHeap.end());
        }
        // Centerdist2.push_back(disHeap);
        // for (int j = 0; j < d; j++)
        // {
        //     pop_heap(disHeap.begin(),disHeap.end());
        //     auto node = disHeap.back();
        //     // cout<<node.dmScore<<endl;
        //     clusterIndex = node.index1;
        //     disHeap.pop_back();
        //     if (cluster[clusterIndex].size()<k){
        //         cluster[clusterIndex].push_back(xi);
        //         break;
        //     }
        // }
    }
    for(int i=0;i<MSize;i++){
        xi = M0[i];
        // min=10e10;
        // clusterIndex=-1;
        // for(int j=0;j<d;j++){
        //     // minTemp = calEuclidean(perturbWorkers[xi.first].loc,muWorkers[j].loc,perturbTasks[xi.second].loc,muTasks[j].loc);
        //     minTemp = Centerdist[i][j];
        //     if(minTemp<min){
        //         // cout<<minTemp<<endl;
        //         min = minTemp;
        //         clusterIndex = j;
        //     }
        // }

        vector<double> tmp = Centerdist[i];
        clusterIndex = min_element(tmp.begin(),tmp.end())-tmp.begin();
        if (clusterIndex>-1)
            cluster[clusterIndex].push_back(xi);

        // disHeap = Centerdist2[i];
        // for (int j = 0; j < d; j++)
        // {
        //     pop_heap(disHeap.begin(),disHeap.end());
        //     auto node = disHeap.back();
        //     clusterIndex = node.index1;
        //     disHeap.pop_back();
        //     if (cluster[clusterIndex].size()<k){
        //         cluster[clusterIndex].push_back(xi);
        //         break;
        //     }
        // }
        
    }
    // 更新簇心
    int ciSize;
    pair<double,double> sumW,sumT;
    pair<int,int> ci;
    pair<double,double> oldWLoc;
    pair<double,double> oldTLoc;
    for(int i=0;i<d;i++){
        ciSize = cluster[i].size();
        sumW.first=0;
        sumW.second=0;
        sumT.first=0;
        sumT.second=0;
        for(int j=0;j<ciSize;j++){
            ci = cluster[i][j];
            sumW.first += perturbWorkers[ci.first].loc.first;
            sumW.second += perturbWorkers[ci.first].loc.second;
            sumT.first += perturbTasks[ci.second].loc.first;
            sumT.second += perturbTasks[ci.second].loc.second;
        }
        if(ciSize !=0){
            sumW.first /= ciSize;
            sumW.second /= ciSize;
            sumT.first /= ciSize;
            sumT.second /= ciSize;
        }        
        oldWLoc = muWorkers[i].loc;
        oldTLoc = muTasks[i].loc;
        
        if (isEqual(sumW.first,oldWLoc.first)&&isEqual(sumW.second,oldWLoc.second)){
            muWorkers[i].workerFlag = true;
        }else{
            muWorkers[i].loc = sumW;
        }
        if (isEqual(sumT.first,oldTLoc.first)&&isEqual(sumT.second,oldTLoc.second)){
            muTasks[i].taskFlag = true;
        }else{
            muTasks[i].loc = sumT;
        }
    }
    //检查质心是否还在更新
    // printf("%.7lf,%.7lf,%.7lf,%.7lf\n",muWorkers[133].loc.first,muWorkers[133].loc.second,muTasks[133].loc.first,muTasks[133].loc.second);
    // cout<<"w:"<<muWorkers[2].loc.first<<"\t"<<muWorkers[2].loc.second<<"\t t:"<<muTasks[2].loc.first<<"\t "<<muTasks[2].loc.second<<endl;
    for (int i = 0; i < d; i++)
    {
        if(!muWorkers[i].workerFlag || !muTasks[i].taskFlag){
            // cout<<"update:"<<cluster[i].size()<<endl;
            return 1;
        }
    }
    return 0;
}

void Kmeans(int&k,vector<vector<pair<int,int>>>&cluster){
    int d;
    int MSize = M0.size();
    if (MSize%k==0)
        d = MSize/k;
    else
        d = MSize/k+1;
    vector<Center> muWorkers,muTasks;
    // calSMatrix();//计算方差矩阵
    //初始化聚类中心点
    initCenter(MSize,d,muWorkers,muTasks);
    cluster.resize(d);
    while(calMeans(MSize,d,muWorkers,muTasks,cluster));
    // for (int i = 0; i < d; i++)
    // {
    //     cout<<cluster[i].size()<<endl;
    // }
    
}

int HEprotocol2(vector<vector<pair<int,int>>>&cluster,vector<int>&failWorkers,vector<int>&failTasks,vector<pair<int,int>>&Mg){
     //HE protocol 
    vector<pair<int,int>>g;
    int clustersize,gsize,successNum=0;
    clustersize = cluster.size();
    for(int i=0;i<clustersize;i++){
        g = cluster[i];
        gsize = g.size();
        successNum += g_M(gsize,g,Mg,failWorkers,failTasks);
    }
    // cout<<"success:"<<successNum<<endl;
    return successNum;
}

int Kmeans_main(){
    // int d;
    // int MSize = M0.size();
    // if (MSize%k==0)
    //     d = MSize/k;
    // else
    //     d = MSize/k+1;
    // vector<Center> muWorkers,muTasks;
    // //初始化聚类中心点
    // initCenter(MSize,d,muWorkers,muTasks);
    
    
    //分类为哪一个簇
    vector<vector<pair<int,int>>> cluster;
    vector<pair<int,int>> Mg;
    vector<int> failWorkers,failTasks;
    int totalSuccess=0;
    /***
    for(int i=0;i<MSize;i++){
        pair<int,int> xi = M0[i];
        double min=10e10,minTemp;
        int clusterIndex=-1;
        for(int j=0;j<d;j++){
            minTemp = calEuclidean(perturbWorkers[xi.first].loc,muWorkers[j].loc,perturbTasks[xi.second].loc,muTasks[j].loc);
            if(minTemp<min){
                // cout<<minTemp<<endl;
                min = minTemp;
                clusterIndex = j;
            }
        }
        cluster[clusterIndex].push_back(xi);
    }
    // 更新簇心
    for(int i=0;i<d;i++){
        int ciSize = cluster[i].size();
        pair<double,double> sumW,sumT;
        sumW.first=0;
        sumW.second=0;
        sumT.first=0;
        sumT.second=0;
        for(int j=0;j<ciSize;j++){
            pair<int,int> ci = cluster[i][j];
            sumW.first += perturbWorkers[ci.first].loc.first;
            sumW.second += perturbWorkers[ci.first].loc.second;
            sumT.first += perturbTasks[ci.second].loc.first;
            sumT.second += perturbTasks[ci.second].loc.second;
        }
        sumW.first /= ciSize;
        sumW.second /= ciSize;
        sumT.first /= ciSize;
        sumT.second /= ciSize;
        
        pair<double,double> oldWLoc = muWorkers[i].loc;
        pair<double,double> oldTLoc = muTasks[i].loc;
        if (sumW != oldWLoc)
            muWorkers[i].loc = sumW;
        else
            muWorkers[i].workerFlag=true;
        if (sumT != oldTLoc)
            muTasks[i].loc = sumT;
        else
            muTasks[i].taskFlag=true;

    }
    **/

    // while(calMeans(MSize,d,muWorkers,muTasks,cluster));
    for (int i = 0; i < lamda; i++)
    {
        cluster.clear();
        Kmeans(k,cluster);
        totalSuccess+=HEprotocol2(cluster,failWorkers,failTasks,Mg);
        // cout<<"fail1:"<<failWorkers.size()<<" "<<failTasks.size()<<endl;
        // cout<<oblivious_M(failWorkers,failTasks)<<endl;
        if(oblivious_Ma(failWorkers,failTasks)==0) break;
        // cout<<"fail2:"<<failWorkers.size()<<" "<<failTasks.size()<<endl;
    }
    // cout<<totalSuccess<<endl;
    // cout<<Mg.size()<<endl;
    
    return totalSuccess;
}